package com.gitee.fastmybatis.spring.boot.autoconfigure.multi;

import com.gitee.fastmybatis.core.FastmybatisConfig;
import com.gitee.fastmybatis.core.MybatisContext;
import com.gitee.fastmybatis.core.ext.SqlSessionFactoryBeanExt;
import com.gitee.fastmybatis.core.ext.spi.SpiContext;
import com.gitee.fastmybatis.core.ext.transaction.AliasTransactionTemplate;
import com.gitee.fastmybatis.core.ext.transaction.TransactionTemplateManager;
import com.gitee.fastmybatis.core.support.plugin.SqlFormatterPlugin;
import com.gitee.fastmybatis.core.util.ClassUtil;
import com.gitee.fastmybatis.spring.boot.autoconfigure.MybatisProperties;
import com.gitee.fastmybatis.spring.boot.autoconfigure.SpringBootVFS;
import com.gitee.fastmybatis.spring.boot.autoconfigure.util.FillUtil;
import org.apache.ibatis.annotations.Mapper;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.session.ExecutorType;
import org.apache.ibatis.session.SqlSessionFactory;
import org.mybatis.spring.SqlSessionTemplate;
import org.mybatis.spring.mapper.ClassPathMapperScanner;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.BeanFactory;
import org.springframework.beans.factory.BeanFactoryAware;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.boot.context.properties.bind.BindResult;
import org.springframework.boot.context.properties.bind.Binder;
import org.springframework.context.ApplicationContext;
import org.springframework.context.EnvironmentAware;
import org.springframework.context.ResourceLoaderAware;
import org.springframework.context.annotation.ImportBeanDefinitionRegistrar;
import org.springframework.core.env.Environment;
import org.springframework.core.io.ResourceLoader;
import org.springframework.core.type.AnnotationMetadata;
import org.springframework.jdbc.datasource.DataSourceTransactionManager;
import org.springframework.util.ObjectUtils;
import org.springframework.util.StringUtils;

import javax.sql.DataSource;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;

public class MapperScannerRegistrar implements
        BeanFactoryAware,
        ResourceLoaderAware,
        EnvironmentAware,
        ImportBeanDefinitionRegistrar {
    private static final Logger logger = LoggerFactory.getLogger(MapperScannerRegistrar.class);

    private BeanFactory beanFactory;
    private ResourceLoader resourceLoader;
    private Environment environment;
    private BeanDefinitionRegistry beanDefinitionRegistry;

    @Override
    public void setResourceLoader(ResourceLoader resourceLoader) {
        this.resourceLoader = resourceLoader;
    }

    @Override
    public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
        this.beanFactory = beanFactory;
    }

    @Override
    public void setEnvironment(Environment environment) {
        this.environment = environment;
    }

    @Override
    public void registerBeanDefinitions(AnnotationMetadata importingClassMetadata, BeanDefinitionRegistry registry) {
        this.beanDefinitionRegistry = registry;

        this.doRegisterBeanDefinitions(registry);
    }

    private void doRegisterBeanDefinitions(BeanDefinitionRegistry registry) {
        MultiDataSourceProps multiDataSourceProps = buildMultiDataSourceProps(environment);
        for (Map.Entry<String, DataSourceProp> entry : multiDataSourceProps.getMultiDatasource().entrySet()) {
            String dsAlias = entry.getKey();
            DataSourceProp prop = entry.getValue();

            DatasourcePropContext datasourcePropContext = DatasourcePropContext.of(entry);
            MyBatisComponent myBatisComponent = buildMyBatisComponent(dsAlias, datasourcePropContext);
            cacheComponent(myBatisComponent);
            MultiDataSourceContext.setMyBatisComponent(dsAlias, myBatisComponent);

            String basePackage = prop.getMybatis().getBasePackage();
            logger.debug("Searching for mappers annotated with @Mapper, package is {}", basePackage);
            Set<String> packages = new HashSet<>();

            ClassPathMapperScanner scanner = new ClassPathMapperScanner(registry);
            if (this.resourceLoader != null) {
                scanner.setResourceLoader(this.resourceLoader);
            }
            //scanner.setSqlSessionFactory(myBatisComponent.getSqlSessionFactory());
            scanner.setSqlSessionTemplate(myBatisComponent.getSqlSessionTemplate());

            if (StringUtils.hasText(basePackage)) {
                String[] packageArr = basePackage.split(";|,");
                packages.addAll(Arrays.asList(packageArr));
            }
            if (logger.isDebugEnabled()) {
                for (String pkg : packages) {
                    logger.debug("Using auto-configuration base package '{}'", pkg);
                }
            }

            scanner.setMarkerInterface(com.gitee.fastmybatis.core.mapper.Mapper.class);
            scanner.setAnnotationClass(Mapper.class);
            scanner.registerFilters();
            scanner.doScan(StringUtils.toStringArray(packages));
        }
    }

    private MyBatisComponent buildMyBatisComponent(String dsAlias, DatasourcePropContext datasourcePropContext) {
        DataSourceProp dataSourceProp = datasourcePropContext.getDataSourceProp();
        checkDatasource(dsAlias, dataSourceProp);

        logger.info("初始化数据源:{}, url={}", dsAlias, dataSourceProp.getUrl());
        DataSource dataSource = DataSourceBuilder.builder(dataSourceProp).build();
        SqlSessionFactory sqlSessionFactory = buildSqlSessionFactory(dataSourceProp, dataSource, null);
        SqlSessionTemplate sqlSessionTemplate = buildSqlSessionTemplate(dataSourceProp, sqlSessionFactory);
        DataSourceTransactionManager transactionManager = new DataSourceTransactionManager(dataSource);
        AliasTransactionTemplate transactionTemplate = new AliasTransactionTemplate(dsAlias, transactionManager);

        return new MyBatisComponent(
                datasourcePropContext,
                dataSource,
                transactionManager,
                sqlSessionFactory,
                sqlSessionTemplate,
                transactionTemplate
        );
    }

    private void cacheComponent(MyBatisComponent component) {
        DatasourcePropContext datasourcePropContext = component.getDatasourcePropContext();
        MybatisContext.addSqlSessionFactory(datasourcePropContext.getDsAlias(),
                component.getSqlSessionFactory(), datasourcePropContext.isMain());

        TransactionTemplateManager.addTransactionTemplate(component.getTransactionTemplate());

        if (datasourcePropContext.getDataSourceProp().getMain()) {
            TransactionTemplateManager.addMainTransactionManager(component.getTransactionManager());
        }
    }

    private SqlSessionFactory buildSqlSessionFactory(
            DataSourceProp dataSourceProp,
            DataSource existDataSource,
            ApplicationContext applicationContext
    ) {
        SqlSessionFactoryBeanExt sessionFactoryBeanExt = buildSqlSessionFactoryBean(dataSourceProp, existDataSource, applicationContext);
        try {
            return sessionFactoryBeanExt.getObject();
        } catch (Exception e) {
            logger.error("Create SqlSessionFactory error", e);
            throw new RuntimeException(e);
        }
    }

    private SqlSessionFactoryBeanExt buildSqlSessionFactoryBean(
            DataSourceProp dataSourceProp,
            DataSource existDataSource,
            ApplicationContext applicationContext
    ) {
        MybatisProperties mybatisProperties = dataSourceProp.getMybatis();

        SqlSessionFactoryBeanExt factory = new SqlSessionFactoryBeanExt();

        DataSource dataSource = existDataSource != null
                ? existDataSource
                : DataSourceBuilder.builder(dataSourceProp).build();
        factory.setDataSource(dataSource);

        factory.setVfs(SpringBootVFS.class);
        if (StringUtils.hasText(mybatisProperties.getConfigLocation())) {
            factory.setConfigLocation(this.resourceLoader.getResource(mybatisProperties.getConfigLocation()));
        }
        org.apache.ibatis.session.Configuration configuration = mybatisProperties.getConfiguration();
        if (configuration == null && !StringUtils.hasText(mybatisProperties.getConfigLocation())) {
            configuration = new org.apache.ibatis.session.Configuration();
        }
        factory.setConfiguration(configuration);
        if (mybatisProperties.getConfigurationProperties() != null) {
            factory.setConfigurationProperties(mybatisProperties.getConfigurationProperties());
        }
        if (StringUtils.hasLength(mybatisProperties.getTypeAliasesPackage())) {
            factory.setTypeAliasesPackage(mybatisProperties.getTypeAliasesPackage());
        }
        if (StringUtils.hasLength(mybatisProperties.getTypeHandlersPackage())) {
            factory.setTypeHandlersPackage(mybatisProperties.getTypeHandlersPackage());
        }
        if (!ObjectUtils.isEmpty(mybatisProperties.resolveMapperLocations())) {
            factory.setMapperLocations(mybatisProperties.resolveMapperLocations());
        }
        List<Class<? extends Interceptor>> pluginsList = mybatisProperties.getPlugins();
        List<Interceptor> interceptors = new ArrayList<>();
        if (!ObjectUtils.isEmpty(pluginsList)) {
            List<Interceptor> interceptorList = pluginsList.stream()
                    .map(aClass ->
                            (Interceptor) ClassUtil.newInstance(aClass))
                    .collect(Collectors.toList());
            interceptors.addAll(interceptorList);
        }
        if (mybatisProperties.isPrintSql()) {
            enableSqlFormatPlugin(interceptors);
        }
        Interceptor[] pluginArr = interceptors.toArray(new Interceptor[0]);
        factory.setPlugins(pluginArr);


        // ====以下是附加属性====
        // dao所在的包名,跟MapperScannerConfigurer的basePackage一致,多个用;隔开
        String basePackage = mybatisProperties.getBasePackage();
        factory.setBasePackage(basePackage);
        factory.setConfig(fastmybatisConfig(mybatisProperties));
        factory.setApplicationContext(applicationContext);
        return factory;
    }

    private void enableSqlFormatPlugin(List<Interceptor> plugins) {
        Optional<Interceptor> first = plugins.stream()
                .filter(data -> data instanceof SqlFormatterPlugin)
                .findFirst();
        if (first.isPresent()) {
            SqlFormatterPlugin interceptor = (SqlFormatterPlugin) first.get();
            interceptor.setEnable(true);
        } else {
            SqlFormatterPlugin interceptor = new SqlFormatterPlugin();
            interceptor.setEnable(true);
            plugins.add(interceptor);
        }
    }

    private FastmybatisConfig fastmybatisConfig(MybatisProperties properties) {
        FastmybatisConfig config = new FastmybatisConfig();
        SpiContext.getBeanExecutor().copyPropertiesIgnoreNull(properties, config);
        if (properties.getFill() != null) {
            config.setFills(FillUtil.buildFills(properties.getFill()));
        }
        return config;
    }

    private SqlSessionTemplate buildSqlSessionTemplate(DataSourceProp dataSourceProp, SqlSessionFactory sqlSessionFactory) {
        ExecutorType executorType = dataSourceProp.getMybatis().getExecutorType();
        if (executorType != null) {
            return new SqlSessionTemplate(sqlSessionFactory, executorType);
        } else {
            return new SqlSessionTemplate(sqlSessionFactory);
        }
    }

    private void checkDatasource(String dbAlias, DataSourceProp dataSourceProp) {
        String prefix = "spring.multi-datasource." + dbAlias + ".";
        checkNotBlank(prefix + "url", dataSourceProp.getUrl());
        checkNotBlank(prefix + "driver-class-name", dataSourceProp.getDriverClassName());
        checkNotBlank(prefix + "username", dataSourceProp.getUsername());
        checkNotBlank(prefix + "password", dataSourceProp.getPassword());
        MybatisProperties mybatisProperties = dataSourceProp.getMybatis();
        checkNotBlank(prefix + "mybatis", mybatisProperties);
        checkNotBlank(prefix + "mybatis.base-package", mybatisProperties.getBasePackage());
    }

    private void checkNotBlank(String name, Object val) {
        if (ObjectUtils.isEmpty(val)) {
            throw new IllegalArgumentException("多数据源配置: '" + name + "' 必填");
        }
    }


    private MultiDataSourceProps buildMultiDataSourceProps(Environment environment) {
        BindResult<MultiDataSourceProps> result = Binder.get(environment).bind("spring", MultiDataSourceProps.class);
        MultiDataSourceProps multiDataSourceProps = result.get();

        int mainCount = 0;
        Set<Map.Entry<String, DataSourceProp>> entries = multiDataSourceProps.getMultiDatasource().entrySet();
        for (Map.Entry<String, DataSourceProp> entry : entries) {
            DataSourceProp prop = entry.getValue();
            Boolean main = prop.getMain();
            if (main) {
                mainCount++;
            }
            if (mainCount > 1) {
                String dsAlias = entry.getKey();
                throw new RuntimeException("只能指定一个主数据源, 参考设置 'spring.multi-datasource." + dsAlias + ".main=false'");
            }
        }

        // 如果没有设置主数据源
        if (mainCount == 0) {
            throw new RuntimeException("必须设置一个主数据源，'spring.multi-datasource.xx.main=true'");
        }

        return multiDataSourceProps;
    }

}
